package cn.edu.flink.tutorial.state;

import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import cn.edu.flink.tutorial.source.UDSourceFunction;

public class KeyedState {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        SingleOutputStreamOperator<Tuple2<Long, Long>> streamSource = env.addSource(new UDSourceFunction()).map(x -> Tuple2.of(x, 1L)).returns(Types.TUPLE(Types.LONG, Types.LONG));


        KeyedStream<Tuple2<Long, Long>, Tuple> keyedOutput = streamSource.keyBy(0);

        DataStream<Tuple2<Long, Long>> sumedOutput = keyedOutput.map(
                new RichMapFunction<Tuple2<Long, Long>, Tuple2<Long, Long>>() {
                    // 键控状态
                    private ValueState<Long> sum;

                    @Override
                    public void open(Configuration parameters) throws Exception {
                        sum = getRuntimeContext().getState(new ValueStateDescriptor<Long>("sum", Long.class, 0L));
                    }

                    @Override
                    public Tuple2<Long, Long> map(Tuple2<Long, Long> integerIntegerTuple2) throws Exception {
                        sum.update(sum.value() + 1);
                        integerIntegerTuple2.f1 = sum.value();
                        return integerIntegerTuple2;
                    }
                }
        );

        sumedOutput.print();

        env.execute();
    }
}
